#include "protocol.h"
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <esp_log.h>

#include "bsp_config.h"
#include <freertos/FreeRTOS.h>
#include <freertos/event_groups.h>

#define TAG "Protocol"
#define PROTOCOL_EVENT_HELLO (1 << 0)

struct protocol
{
    protocol_transport_t *transport;
    text_callback_t stt_callback;
    text_callback_t llm_callback;
    action_callback_t tts_start_callback;
    text_callback_t tts_sentence_start_callback;
    action_callback_t tts_stop_callback;
    action_callback_t audio_channel_close_callback;
    audio_callback_t audio_in_callback;
    text_callback_t iot_callback;

    char *session_id;
    int server_sample_rate;

    EventGroupHandle_t event_group;
};

static void protocol_audio_in_callback(void *arg, void *data, size_t len);
static void protocol_channel_close_callback(void *arg);
static void protocol_text_callback(void *arg, char *text);

static protocol_transport_callbacks_t callbacks = {
    .channel_close_callback = protocol_channel_close_callback,
    .audio_callback = protocol_audio_in_callback,
    .text_callback = protocol_text_callback,
};

// protocol_audio_in_callback、protocol_text_callback、protocol_channel_close_callback
// 三个函数是供协议传输层调用的回调
// 从协议传输层上来的数据，一般两种类型，音频和文本。还有一种可能是服务器关闭了信道连接
static void protocol_audio_in_callback(void *arg, void *data, size_t len)
{
    protocol_t *protocol = arg;
    if (protocol->audio_in_callback)
    {
        protocol->audio_in_callback(data, len);
    }
}

static void protocol_text_callback(void *arg, char *text)
{
    protocol_t *protocol = arg;
    cJSON *root = cJSON_Parse(text);
    if (!root)
    {
        ESP_LOGE(TAG, "Failed to parse JSON");
        ESP_LOGE(TAG, "JSON: %s", text);
        return;
    }

    cJSON *type_obj = cJSON_GetObjectItem(root, "type");
    if (!cJSON_IsString(type_obj))
    {
        ESP_LOGE(TAG, "Failed to parse reply type");
        cJSON_Delete(root);
        return;
    }

    if (strcmp(type_obj->valuestring, "hello") == 0)
    {
        // 解析hello数据
        cJSON *audio_params_obj = cJSON_GetObjectItem(root, "audio_params");
        cJSON *session_id_obj = cJSON_GetObjectItem(root, "session_id");

        if (!audio_params_obj || !cJSON_IsString(session_id_obj))
        {
            ESP_LOGE(TAG, "Invalid hello message, %s", text);
            cJSON_Delete(root);
            return;
        }
        protocol->session_id = strdup(session_id_obj->valuestring);

        cJSON *sample_rate_obj = cJSON_GetObjectItem(audio_params_obj, "sample_rate");
        if (!cJSON_IsNumber(sample_rate_obj))
        {
            ESP_LOGE(TAG, "Invalid hello message, %s", text);
            cJSON_Delete(root);
            return;
        }

        protocol->server_sample_rate = sample_rate_obj->valueint;
        xEventGroupSetBits(protocol->event_group, PROTOCOL_EVENT_HELLO);
    }
    else if (strcmp(type_obj->valuestring, "stt") == 0 && protocol->stt_callback)
    {
        // 解析stt数据
        cJSON *text_obj = cJSON_GetObjectItem(root, "text");
        if (!cJSON_IsString(text_obj))
        {
            ESP_LOGE(TAG, "Invalid stt message, %s", text);
            cJSON_Delete(root);
            return;
        }
        protocol->stt_callback(text_obj->valuestring);
    }
    else if (strcmp(type_obj->valuestring, "llm") == 0 && protocol->llm_callback)
    {
        // 解析llm数据
        cJSON *text_obj = cJSON_GetObjectItem(root, "text");
        if (!cJSON_IsString(text_obj))
        {
            ESP_LOGE(TAG, "Invalid stt message, %s", text);
            cJSON_Delete(root);
            return;
        }
        protocol->llm_callback(text_obj->valuestring);
    }
    else if (strcmp(type_obj->valuestring, "tts") == 0)
    {
        // 查看state
        cJSON *state_obj = cJSON_GetObjectItem(root, "state");
        if (!cJSON_IsString(state_obj))
        {
            ESP_LOGE(TAG, "TTS state is not a string");
            cJSON_Delete(root);
            return;
        }

        if (strcmp(state_obj->valuestring, "start") == 0 && protocol->tts_start_callback)
        {
            protocol->tts_start_callback();
        }
        else if (strcmp(state_obj->valuestring, "stop") == 0 && protocol->tts_stop_callback)
        {
            protocol->tts_stop_callback();
        }
        else if (strcmp(state_obj->valuestring, "sentence_start") == 0 && protocol->tts_sentence_start_callback)
        {
            cJSON *text_obj = cJSON_GetObjectItem(root, "text");
            if (!cJSON_IsString(text_obj))
            {
                ESP_LOGE(TAG, "Invalid stt message, %s", text);
                cJSON_Delete(root);
                return;
            }
            protocol->tts_sentence_start_callback(text_obj->valuestring);
        }
    }
    else if (strcmp(type_obj->valuestring, "iot") == 0)
    {
        cJSON *command_array = cJSON_GetObjectItem(root, "commands");
        if (command_array)
        {
            protocol->iot_callback((char *)command_array);
        }
    }
    cJSON_Delete(root);
}

static void protocol_channel_close_callback(void *arg)
{
    protocol_t *protocol = arg;
    if (protocol->audio_channel_close_callback)
    {
        protocol->audio_channel_close_callback();
    }
}



// 检查能否发送文本
static bool protocol_can_send_text(protocol_t *protocol)
{
    if (!protocol->transport->is_audio_channel_open || !protocol->transport->send_text)
    {
        ESP_LOGE(TAG, "Send text not supported");
        return false;
    }

    if (!protocol->transport->is_audio_channel_open(protocol->transport))
    {
        ESP_LOGE(TAG, "Audio channel not open");
        return false;
    }
    return true;
}

// 协议初始化
protocol_t *protocol_create(protocol_transport_type_t type)
{
    protocol_t *protocol = malloc(sizeof(protocol_t));
    assert(protocol);
    memset(protocol, 0, sizeof(protocol_t));

    protocol->event_group = xEventGroupCreate();
    assert(protocol->event_group);

    switch (type)
    {
    case PROTOCOL_TRANSPORT_TYPE_WEBSOCKET:
        protocol->transport = protocol_transport_create_websocket(&callbacks, protocol);
        break;

    default:
        ESP_LOGE(TAG, "Unsupported protocol transport type");
        break;
    }

    if (protocol->transport->start)
    {
        protocol->transport->start(protocol->transport);
    }

    return protocol;
}



// protocol_open_audio_channel、protocol_close_audio_channel、protocol_is_audio_channel_open
// 负责开启、关闭、检查信道状态
bool protocol_open_audio_channel(protocol_t *protocol)
{
    if (!protocol->transport->open_audio_channel)
    {
        ESP_LOGE(TAG, "Open audio channel not supported");
        return false;
    }

    return protocol->transport->open_audio_channel(protocol->transport);
}

void protocol_close_audio_channel(protocol_t *protocol)
{
    if (!protocol->transport->close_audio_channel)
    {
        ESP_LOGE(TAG, "Close audio channel not supported");
        return;
    }
    protocol->transport->close_audio_channel(protocol->transport);
}

bool protocol_is_audio_channel_open(protocol_t *protocol)
{
    if (protocol->transport->is_audio_channel_open)
    {
        return protocol->transport->is_audio_channel_open(protocol->transport);
    }

    return false;
}


// 一系列send函数，是单片机主动向服务器发送的请求
void protocol_send_audio(protocol_t *protocol, void *data, size_t len)
{
    if (!protocol->transport->is_audio_channel_open || !protocol->transport->send_audio)
    {
        ESP_LOGE(TAG, "Send audio not supported");
        return;
    }

    if (protocol->transport->is_audio_channel_open(protocol->transport))
    {
        protocol->transport->send_audio(protocol->transport, data, len);
    }
    else
    {
        ESP_LOGE(TAG, "Audio channel not open");
    }
}

void protocol_send_hello(protocol_t *protocol)
{
    char *text = NULL;
    char *transport = NULL;

    if (!protocol_can_send_text(protocol))
    {
        return;
    }

    switch (protocol->transport->type)
    {
    case PROTOCOL_TRANSPORT_TYPE_WEBSOCKET:
        transport = "websocket";
        break;

    default:
        transport = "unknown";
        break;
    }
    asprintf(&text, "{\"audio_params\": {\"channels\": %d, \"format\": \"opus\", \"frame_duration\": 60, \"sample_rate\": %d}, \"transport\": \"%s\", \"type\": \"hello\", \"version\": 1}",
             BSP_CODEC_CHANNEL_COUNT,
             BSP_CODEC_SAMPLE_RATE,
             transport);
    protocol->transport->send_text(protocol->transport, text);
    free(text);

    // 等待服务器回复
    xEventGroupWaitBits(protocol->event_group, PROTOCOL_EVENT_HELLO, pdTRUE, pdTRUE, portMAX_DELAY);
}

void protocol_send_wake_word(protocol_t *protocol, char *word)
{
    char *text = NULL;
    if (!protocol_can_send_text(protocol))
    {
        return;
    }

    asprintf(&text, "{\"session_id\": \"%s\", \"state\": \"detect\", \"text\": \"%s\", \"type\": \"listen\"}",
             protocol->session_id ? protocol->session_id : "",
             word);
    protocol->transport->send_text(protocol->transport, text);
    free(text);
}

void protocol_send_start_listening(protocol_t *protocol)
{
    char *text = NULL;
    if (!protocol_can_send_text(protocol))
    {
        return;
    }
    asprintf(&text, "{\"mode\": \"manual\", \"session_id\": \"%s\", \"state\": \"start\", \"type\": \"listen\"}",
             protocol->session_id ? protocol->session_id : "");
    protocol->transport->send_text(protocol->transport, text);
    free(text);
}

void protocol_send_stop_listening(protocol_t *protocol)
{
    char *text = NULL;
    if (!protocol_can_send_text(protocol))
    {
        return;
    }
    asprintf(&text, "{\"session_id\": \"%s\", \"state\": \"stop\", \"type\": \"listen\"}",
             protocol->session_id ? protocol->session_id : "");
    protocol->transport->send_text(protocol->transport, text);
    free(text);
}

void protocol_abort_speaking(protocol_t *protocol, protocol_abort_reason_t reason)
{
    char *text = NULL;
    if (!protocol_can_send_text(protocol))
    {
        return;
    }

    switch (reason)
    {
    case PROTOCOL_ABORT_REASON_WAKEUP:
        asprintf(&text, "{\"reason\": \"wake_word_detected\", \"session_id\": \"%s\", \"type\": \"abort\"}",
                 protocol->session_id ? protocol->session_id : "");
        break;

    default:
        asprintf(&text, "{\"session_id\": \"%s\", \"type\": \"abort\"}",
                 protocol->session_id ? protocol->session_id : "");
        break;
    }

    protocol->transport->send_text(protocol->transport, text);
    free(text);
}

void protocol_send_iot_descriptor(protocol_t *protocol, cJSON *descriptor_json)
{
    if (!protocol_can_send_text(protocol))
    {
        cJSON_Delete(descriptor_json);
        return;
    }

    cJSON *root = cJSON_CreateObject();
    cJSON_AddItemToObject(root, "descriptors", descriptor_json);
    cJSON_AddStringToObject(root, "type", "iot");
    cJSON_AddStringToObject(root, "session_id", protocol->session_id);
    cJSON_AddBoolToObject(root, "update", true);
    char *json_str = cJSON_PrintUnformatted(root);
    cJSON_Delete(root);

    protocol->transport->send_text(protocol->transport, json_str);
    free(json_str);
}

void protocol_send_iot_state(protocol_t *protocol, cJSON *state_json)
{
    if (!protocol_can_send_text(protocol))
    {
        cJSON_Delete(state_json);
        return;
    }

    cJSON *root = cJSON_CreateObject();
    cJSON_AddItemToObject(root, "states", state_json);
    cJSON_AddStringToObject(root, "type", "iot");
    cJSON_AddStringToObject(root, "session_id", protocol->session_id);
    cJSON_AddBoolToObject(root, "update", true);
    char *json_str = cJSON_PrintUnformatted(root);
    cJSON_Delete(root);

    protocol->transport->send_text(protocol->transport, json_str);
    free(json_str);
}



// callback系列函数，主要负责处理来自于服务器的命令所产生的回调
void protocol_set_stt_callback(protocol_t *protocol, text_callback_t callback)
{
    protocol->stt_callback = callback;
}

void protocol_set_llm_callback(protocol_t *protocol, text_callback_t callback)
{
    protocol->llm_callback = callback;
}

void protocol_set_tts_start_callback(protocol_t *protocol, action_callback_t callback)
{
    protocol->tts_start_callback = callback;
}

void protocol_set_tts_stop_callback(protocol_t *protocol, action_callback_t callback)
{
    protocol->tts_stop_callback = callback;
}

void protocol_set_tts_sentence_start_callback(protocol_t *protocol, text_callback_t callback)
{
    protocol->tts_sentence_start_callback = callback;
}

void protocol_set_iot_callback(protocol_t *protocol, text_callback_t callback)
{
    protocol->iot_callback = callback;
}

void protocol_set_audio_in_callback(protocol_t *protocol, audio_callback_t callback)
{
    protocol->audio_in_callback = callback;
}

void protocol_set_audio_channel_close_callback(protocol_t *protocol, action_callback_t callback)
{
    protocol->audio_channel_close_callback = callback;
}

